package org.tigr.htc.common;

/*
This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.Socket;
import java.net.UnknownHostException;

import org.apache.log4j.Logger;

import com.mindbright.jca.security.SecureRandom;
import com.mindbright.ssh2.SSH2ConsoleRemote;
import com.mindbright.ssh2.SSH2SimpleClient;
import com.mindbright.ssh2.SSH2Transport;
import com.mindbright.util.RandomSeed;
import com.mindbright.util.SecureRandomAndPad;

/**
 * The <b>SSHCommand </b> class is uses a SSH client to connect
 * to the given machine, run a command and grab its output.
 */
public class SSHCommand {
    static Logger log = Logger.getLogger(SSHCommand.class);
    final static int DEFAULT_PORT = 22;
    final static int MAX_CONS = 10;
    
    protected static FIFOSemaphore fs = new FIFOSemaphore(MAX_CONS);
    
    protected String hostname;
    protected int port;
    protected String user;
    protected String passwd;
    protected StringBuffer output;
    //    protected String error;
    
    public SSHCommand(String p_hostname, int p_port, String p_user, String p_passwd) {
        hostname = p_hostname;
        port = p_port;
        user = p_user;
        passwd = p_passwd;
    }
    
    public SSHCommand(String p_hostname, String p_user, String p_passwd) {
        hostname = p_hostname;
        port = DEFAULT_PORT;
        user = p_user;
        passwd = p_passwd;
    }
    
    public SSHCommand(String p_hostname, String p_passwd) {
        hostname = p_hostname;
        port = DEFAULT_PORT;
        user = System.getProperty("user.name");
        passwd = p_passwd;
    }
    
    public SSHCommand(int p_port, String p_user, String p_passwd) throws UnknownHostException {
        hostname = InetAddress.getLocalHost().getHostName();
        port = p_port;
        user = p_user;
        passwd = p_passwd;
    }
    
    public SSHCommand(String p_passwd) throws UnknownHostException {
        hostname = InetAddress.getLocalHost().getHostName();
        port = DEFAULT_PORT;
        user = System.getProperty("user.name");
        passwd = p_passwd;
    }
    
    public int runSudoCommand(String p_cmd, String p_user) {
        OutputStream stdin = null;
        boolean result = false;
        
        try {
            // reset output buffer
            output = new StringBuffer();
            /*
             * Connect to the server and authenticate using plain password
             * authentication (if other authentication method needed check other
             * constructors for SSH2SimpleClient).
             */
            log.debug(" set connection to " + hostname + " and " + port);
            Socket serverSocket     = new Socket(hostname, port);
            SSH2Transport transport = new SSH2Transport(serverSocket,createSecureRandom());
            log.debug(" using user " + user);
            SSH2SimpleClient client = new SSH2SimpleClient(transport, user, passwd);
            
            /*
             * Create the remote console to use for command execution.
             */
            SSH2ConsoleRemote console =
                new SSH2ConsoleRemote(client.getConnection(), null, System.err);
            
            /*
             * Run the command (returns a boolean indicating success, we ignore
             * it here). Here we redirect stdout and stderr of the remote command
             * execution to our own stdout and stderr for simplicity.
             */
            log.debug(" running sudo -k to clear cache");
            if(!console.command("sudo -k")) {
                log.debug(" cache clear didn't succeed");
            }
            
            stdin = console.getStdIn();
            log.debug(" running sudo command " + "echo passwd  | sudo -S -u " + p_user + " " + p_cmd);
            if(!console.command("echo \"" + passwd + "\" | sudo -S -u " + p_user + " " + p_cmd)) {
                log.debug("Failed: sudo -S -u " + p_user + " " + p_cmd);
            }
            // send passwd across
            //	    	    log.debug(" send passwd");
            //	    	    stdin.write((passwd + "\n").getBytes());
            //	    	    log.debug(" flush ");
            //		    stdin.flush();
            //		    stdin.close();
            
            /*
             * Fetch the internal stdout stream and wrap it in a BufferedReader
             * for convenience.
             */
            BufferedReader stdout =
                new BufferedReader(new InputStreamReader(console.getStdOut()));
            
            /*
             * Read all output sent to stdout (line by line) and append to output buffer
             * 
             */
            String line;
            while((line = stdout.readLine()) != null) {
                output.append(line).append("\n");
            }
            
            /*
             * Retrieve the exit status of the command (from the remote end).
             */
            int exitStatus = console.waitForExitStatus();
            
            
            /*
             * Disconnect the transport layer gracefully
             */
            transport.normalDisconnect("User disconnects");
            
            /*
             * Exit with same status as remote command did
             */
            return(exitStatus);
            
        } catch (Exception e) {
            log.debug("An error occured: " + e, e);
            return 127;
        }
        
    }
    
    /**
     * run a command as user on a machine.
     */
    public int runCommand(String p_cmd) {
        
        
        boolean result = false;
        
        try {
            // try to acquire a connection
            log.debug("try to acquire semaphore with count " + fs.getCount());
            fs.acquire();
            log.debug(" acquired a semaphore for ssh connection");
            
            // reset output buffer
            output = new StringBuffer();
            /*
             * Connect to the server and authenticate using plain password
             * authentication (if other authentication method needed check other
             * constructors for SSH2SimpleClient).
             */
            log.debug("Set connection to " + hostname + " and " + port);
            Socket serverSocket     = new Socket(hostname, port);
            SSH2Transport transport = new SSH2Transport(serverSocket,createSecureRandom());
            log.debug(" using user " + user);
            SSH2SimpleClient client = new SSH2SimpleClient(transport, user, passwd);
            
            /*
             * Create the remote console to use for command execution.
             */
            SSH2ConsoleRemote console =
                new SSH2ConsoleRemote(client.getConnection());
            
            /*
             * Run the command (returns a boolean indicating success, we ignore
             * it here). Here we redirect stdout and stderr of the remote command
             * execution to our own stdout and stderr for simplicity.
             */
            log.debug(" running " + p_cmd);
            result = console.command(p_cmd);
            
            /*
             * Fetch the internal stdout stream and wrap it in a BufferedReader
             * for convenience.
             */
            BufferedReader stdout =
                new BufferedReader(new InputStreamReader(console.getStdOut()));
            
            /*
             * Read all output sent to stdout (line by line) and print it to our
             * own stdout.
             */
            String line;
            while((line = stdout.readLine()) != null) {
                output.append(line).append("\n");
            }
            
            /*
             * Retrieve the exit status of the command (from the remote end).
             */
            int exitStatus = console.waitForExitStatus();
            
            
            /*
             * Disconnect the transport layer gracefully
             */
            transport.normalDisconnect("User disconnects");
            
            /*
             * Exit with same status as remote command did
             */
            return(exitStatus);
            
        } catch(InterruptedException e){
            log.warn(" thread was interrupted during a semaphore wait" + e, e);
            return 127;
        } catch (Exception e) {
            log.error("An error occured: " + e.getMessage(), e);
            return 127;
        } finally {
            log.debug("releasing semphore");
            fs.release();
            log.debug(" released has count" + fs.getCount());
        }
    }
    
    public String getOutput() {
        return output.toString();
    }
    
    public static SecureRandomAndPad createSecureRandom() {
        /* NOTE, this is how it should be done if you want good
         * randomness, however good randomness takes time so we settle with
         * just some low-entropy garbage here.
         *
         * RandomSeed seed = new RandomSeed("/dev/random", "/dev/urandom");
         * byte[] s = seed.getBytesBlocking(20);
         * return new SecureRandomAndPad(new SecureRandom(s));
         *
         */
        byte[] seed = RandomSeed.getSystemStateHash();
        return new SecureRandomAndPad(new SecureRandom(seed));
    }
    
    /**
     * For testing only, not for general use
     * 
     *
     */
    public static void main(String args[]) {
        String machine[] = "fake,fake2,localhost".split(",");
        String PASS = "redgun3";
        String submitCmd = "/bin/ls -l";
        String luser = "sommer";
        
        SSHCommand sshCmd = new SSHCommand(machine[0], PASS);
        int exitValue = sshCmd.runSudoCommand(submitCmd, luser);
        System.out.println("Return value: " + exitValue);
        String line = sshCmd.getOutput();
        System.out.println(" condor_submit result " + line);
        
    }
    
    
}
